/*
    Copyright 2006-2012 Patrik Jonsson, sunrise@familjenjonsson.org

    This file is part of Sunrise.

    Sunrise is free software; you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation; either version 3 of the License, or
    (at your option) any later version.

    Sunrise is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with Sunrise.  If not, see <http://www.gnu.org/licenses/>.

*/

/// \file
/// Definitions of the polychromatic_grid functions.

#include "polychromatic_grid.h"
#include "grid_factory.h"
#include "grid-fits.h"
#include "CCfits/CCfits"
#include "blitz-fits.h"
#include "vecops.h"
#include <vector>
#include "density_generator.h"
#include "arepo_grid.h"
#include "emission.h"
#include "arepo_emission.h"

using namespace CCfits;
using namespace std;
using namespace blitz;


template<typename T>
mcrx::polychromatic_grid<T>::polychromatic_grid (boost::shared_ptr<T_grid_impl> g) :
  g_(g)
{
  // can't call make_blocks here because there's no data in the
  // grid... it'll have to be called from derived constructors.
}

template<typename T>
mcrx::polychromatic_grid<T>::polychromatic_grid (boost::shared_ptr<T_grid_impl> g,
						 CCfits::ExtHDU& data_hdu,
						 const density_generator& gen,
						 int n_lambda) :
  g_(g)
{
  stupid(data_hdu, gen, n_lambda);
}


template<typename T>
void
mcrx::polychromatic_grid<T>::stupid (CCfits::ExtHDU& data_hdu,
				     const density_generator& gen,
				     int n_lambda)
{
  const int this_task = mpi_rank();

  // we cast the grid pointer to an adaptive_grid here since this is
  // only called for that type and the domain decomposition interface
  // is not common to all types.
  adaptive_grid<typename T_grid_impl::T_data>* dummy;

  adaptive_grid<typename T_grid_impl::T_data>& gref =
    dynamic_cast<adaptive_grid<typename T_grid_impl::T_data>&>(*g_);

  const pair<size_t, size_t> domain = gref.domain_index();

  cout << "Task " << mpi_rank() << " reading grid cell data for "
       << domain.second << " cells starting at " << domain.first << endl;

  Column& c_m_g = data_hdu.column("mass_gas");
  Column& c_m_metals = data_hdu.column("mass_metals");

  std::vector<T_float> m_g, m_metals;
  std::vector<vec3d> p_g;
  c_m_g.read(m_g, domain.first+1, domain.first+domain.second);
  c_m_metals.read(m_metals, domain.first+1, domain.first+domain.second);

  T_float momcon=1;
  try {
    Column& c_p_g = data_hdu.column("p_gas");
    read(c_p_g, p_g, domain.first+1, domain.second);

    // we need to convert momentum, which is in units of
    // mass*length/time, to units of mass*m/s (since what we want is the
    // velocity)
    momcon = units::convert(c_p_g.unit(),c_m_g.unit()+"*m/s");
  }
  catch (Table::NoSuchColumn&) {
    // no momentum. set to zero
    p_g.assign(m_g.size(), vec3d(0,0,0));
  }

  assert(m_g.size()==domain.second);
  assert(m_metals.size()==domain.second);
  assert(p_g.size()==domain.second);

  int i = 0;
  for (iterator c = begin (), e=end(); c != e; ++c, ++i) {
    // make an absorber with the appropriate densities. We only know
    // density of metals and gas so we ask the density generator to
    // give us a density vector, which we put into our big block.
    assert(c.volume()>0);
    assert(m_g [i]>=0);
    assert(m_metals [i]>=0);
    T_densities rho (gen.make_density_vector(m_g [i]/c.volume(),
					     m_metals [i]/c.volume()));
    ASSERT_ALL(rho>=0);
    T_lambda intens (n_lambda);
    intens=0;
    vec3d vel(p_g[i]*momcon/m_g[i]);
    if(m_g[i]==0) {
      vel=0.0;
    }

    // velocity better be smaller than 10000km/s.
    assert(mag(vel)<1e7);
    T_data d(rho, vel, intens);
    DEBUG(3,cout << "Cell velocity " << d.velocity() << " in polychromatic_grid.cc" <<  endl;);

    // if the cell has no data member, default-construct one before assigning
    if(!c->data()) {
      c->set_data(new typename T_grid_impl::T_data() );
    }
    c->data()->get_absorber() = d;
  }
  assert(i==domain.second);

  // and now consolidate into continuous blocks
  make_blocks();
}


/** This function goes through the grid and consolidates the array
    members in the cell data objects to point to be subarrays of one
    contiguous array. */
template<typename T>
void mcrx::polychromatic_grid<T>::make_blocks()
{
  const int nc=n_cells();
  const int this_task = mpi_rank();

  // We also want to allocate blocks for the densities and intensities
  // arrays of the cells. For this we need to know the number of
  // species and wavelengths.

  const int n_species = begin()->data()->get_absorber().densities().size();
  const int n_lambda = begin()->data()->get_absorber().intensity().size();

  cout << "Allocating a (" << nc << ", " << n_species
       << ") memory block for cell densities, "
       << nc*n_species*sizeof(T_float)*1.0/(1024*1024*1024) << " GB.\n";
  rho_.resize(nc, n_species);
  cout << "Allocating a (" << nc << ", " << n_lambda
       << ") memory block for cell intensities, "
       << nc*n_lambda*sizeof(T_float)*1.0/(1024*1024*1024) << " GB.\n";
  intensities_.resize(nc, n_lambda);
  intensities_=0.;

  // go through and make the arrays reference the master array
  int i = 0;
  for (iterator c = begin (), e=end(); c != e; ++c, ++i) {
    assert(c->data());

    T_densities rho;
    rho.weakReference(rho_(i,Range::all()));
    rho = c->data()->get_absorber().densities();
    c->data()->get_absorber().densities().reference(rho);

    T_lambda intens;
    if(n_lambda>0) {
      intens.weakReference(intensities_(i,Range::all()));
      intens = c->data()->get_absorber().intensity();
      c->data()->get_absorber().intensity().reference(intens);
    }
  }
  assert(i==nc);

  ASSERT_ALL(intensities_==intensities_);
  ASSERT_ALL(intensities_< 1e300);
}



/** Destructor.  The destructor must check if we have used placement
    new to allocate the cells and explicitly call the destructors for
    all the cell data if that's the case.  (Note that the destructor
    for the mcrx::adaptive_grid class does the same for the subgrids
    and the grid cells, but the cell data is stored separately.) */
template<typename T>
mcrx::polychromatic_grid<T>::~polychromatic_grid ()
{
  // The array blocks will selfdestruct implicitly and this is safe
  // because at this point, the weak references in the absorber
  // objects have already been destroyed.
}


namespace mcrx {
// Explicitly intantiate the templates necessary to avoid having to
// include everything
template class polychromatic_grid<
  adaptive_grid<
    cell_data<
      typename emitter_types<adaptive_grid, 
			     polychromatic_policy, local_random>::T_emitter,
      absorber<array_1> > > >;

template class polychromatic_grid<adaptive_grid<absorber<array_1> > >;

#ifdef WITH_AREPO
template class polychromatic_grid<
  arepo_grid<
    cell_data<
      typename emitter_types<arepo_grid, 
			     polychromatic_policy, local_random>::T_emitter,
      absorber<array_1> > > >;
#endif
}
